#该部分主要功能是
#选择一段时间，图上呈现该段时间内的上下车热点图像
#进而可以针对不同时间段（如工作日或节假日）进行上下客热点区域/路线和城市规划的分析
import pandas as pd
import numpy as np
from sklearn.cluster import KMeans

def get_sub_trajectory(df1):# 对某辆出租车轨迹进行分段，提取上下车点
    loads=[]
    no_loads=[]
    on_board=[]
    pick_up=[]
    drop_off=[]
    # 辅助记录
    idx1=-1 #记录每一段轨迹的开始
    idx2=-1 #记录每一段轨迹的结束
    old_status=''
    for index, row in df1.iterrows():
        status=row['empty/load']
        # 初始化
        if index==0:
            idx1=index
            old_status=status
        # 判断状态是否转变
        # 记录状态改变的行索引
        if status!=old_status:
            sub_df=df1[idx1:idx2+1]
            if old_status=='重车':
                loads.append(sub_df)
                drop_off.append((row['lon'],row['lat']))
            else:
                no_loads.append(sub_df)
                pick_up.append((row['lon'],row['lat']))
            idx1=index
            idx2=index
            old_status=status
        else:
            idx2=index
    sub_df=df1[idx1:idx2+1]
    if old_status=='重车':
        loads.append(sub_df)
    else:
        no_loads.append(sub_df)
    return loads,no_loads,pick_up,drop_off

def getPickAndDrop(df):
    # 对所有的出租车
    pick=pd.DataFrame()
    drop=pd.DataFrame()
    for tid in set(df['id']):
        df1=df[df['id']==tid].reset_index(drop=True)
        x,x,pick_up,drop_off=get_sub_trajectory(df1)
        if len(pick_up)>0:
            pick_df=pd.DataFrame(pick_up)
            pick_df.columns = ['lon','lat']
            pick_df['PickorDrop'] = 'PICK'
            pick=pick.append(pick_df)
        if len(drop_off)>0:
            drop_df=pd.DataFrame(drop_off)
            drop_df.columns = ['lon','lat']
            drop_df['PickorDrop'] = 'DORP'
            drop=drop.append(drop_df)
    return pick,drop

